#include "capture.h"

static int _analyze_tcp_user(Packet_t* packet)
{
    return 1;
}
static int _analyze_udp_user(Packet_t* packet)
{
    struct udphdr* udp_hdr = packet->udp_hdr;
    //DNS
    if (ntohs(udp_hdr->dest) == 53 || ntohs(udp_hdr->source) == 53) {
        packet->dns_hdr = packet->data;
    }

    return 1;
}

static int _analyze_user_data(Packet_t* packet)
{
    struct udphdr* udp_hdr = packet->udp_hdr;
    struct tcphdr* tcp_hdr = packet->tcp_hdr;
    if (udp_hdr) {
        _analyze_udp_user(packet);
    }
    if (tcp_hdr) {
        _analyze_tcp_user(packet);
    }
    return 1;
}

static int _analyze_icmp(Packet_t* packet)
{
    uint8_t* next_hdr = (uint8_t*)packet->icmp_hdr + sizeof(struct icmp);
    packet->data = (uint8_t*)next_hdr;
    return 1;
}
static int _analyze_udp(Packet_t* packet)
{
    uint8_t* next_hdr = (uint8_t*)packet->udp_hdr + sizeof(struct udphdr);
    packet->data = (uint8_t*)next_hdr;
    _analyze_user_data(packet);
    return 1;
}
static int _analyze_tcp(Packet_t* packet)
{
    uint8_t* next_hdr = (uint8_t*)packet->tcp_hdr + sizeof(struct tcphdr);
    packet->data = (uint8_t*)next_hdr;
    _analyze_user_data(packet);
    return 1;
}

static int _analyze_ip(Packet_t* packet)
{
    uint8_t* next_hdr = (uint8_t*)packet->ip_hdr + (packet->ip_hdr->ihl * 4);
    switch (packet->ip_hdr->protocol) {
    case IPPROTO_TCP:
        packet->tcp_hdr = (struct tcphdr*)next_hdr;
        _analyze_tcp(packet);
        break;
    case IPPROTO_UDP:
        packet->udp_hdr = (struct udphdr*)next_hdr;
        _analyze_udp(packet);
        break;
    case IPPROTO_ICMP:
        packet->icmp_hdr = (struct icmp*)next_hdr;
        _analyze_icmp(packet);
        break;
    }
    return 1;
}

static int _analyze_ppp_ses(Packet_t* packet)
{

    uint8_t* next_hdr = (uint8_t*)packet->ppp_hdr + sizeof(struct ppphdr);
    switch (ntohs(packet->ppp_hdr->proto)) {
    case ETH_P_IP:
        packet->ip_hdr = (struct iphdr*)next_hdr;
        _analyze_ip(packet);
        break;
    default:
        break;
    }
    return 1;
}

static int _analyze_vlan(Packet_t* packet)
{

    uint8_t* next_hdr = (uint8_t*)packet->vlan_hdr + sizeof(struct vlanhdr);
    switch (ntohs(packet->vlan_hdr->type)) {
    case ETH_P_IP:
        packet->ip_hdr = (struct iphdr*)next_hdr;
        _analyze_ip(packet);
        break;
    case ETH_P_PPP_SES:
        packet->ppp_hdr = (struct ppphdr *)next_hdr;
        _analyze_ppp_ses(packet);
        break;
    default:
        break;
    }
    return 1;
}

static int _analyze_packet(Packet_t* packet)
{
    packet->eth_hdr = (struct ethhdr*)packet->buf;
    uint8_t* next_hdr = (uint8_t*)packet->eth_hdr + sizeof(struct ethhdr);
    switch (ntohs(packet->eth_hdr->h_proto)) {
    case ETH_P_IP:
        packet->ip_hdr = (struct iphdr*)next_hdr;
        _analyze_ip(packet);
        break;
    case ETH_P_8021Q:
        packet->vlan_hdr = (struct vlanhdr*)next_hdr;
        _analyze_vlan(packet);
        break;
    default:
        break;
    }
    return 1;
}

static void* Analyze(void* arg)
{
    Packet_t* packet = (Packet_t*)arg;
    _analyze_packet(packet);
    Hook(packet);
    free(packet);
}

int capture_bind(Capture_t* cap, char* ifName)
{
    struct sockaddr_ll sl;
    struct ifreq ifr;

    memset(&sl, 0x00, sizeof(sl));
    memset(&ifr, 0x00, sizeof(ifr));

    sl.sll_family = AF_PACKET;
    sl.sll_protocol = htons(ETH_P_ALL);

    strncpy(ifr.ifr_name, ifName, sizeof(ifr.ifr_name));
    ioctl(cap->raw_fd, SIOCGIFINDEX, &ifr);
    sl.sll_ifindex = ifr.ifr_ifindex;
    if (0 > bind(cap->raw_fd, (struct sockaddr*)&sl, sizeof(sl))) {
        log_syserr("cap bind error.");
        return -1;
    }
    return 1;
}
int capture_set_promisc(Capture_t* cap, char* ifName)
{
    struct ifreq ifr;
    strcpy(ifr.ifr_name, ifName);

    if (ioctl(cap->raw_fd, SIOCGIFFLAGS, &ifr) == -1) {
        log_syserr("ioctl error.");
        return -1;
    }

    ifr.ifr_flags |= IFF_PROMISC;
    if (ioctl(cap->raw_fd, SIOCSIFFLAGS, &ifr) == -1) {
        log_syserr("ioctl error.");
        return -1;
    }

    return 0;
}

Capture_t* capture_new()
{
    Capture_t* cap;

    cap = (Capture_t*)malloc(sizeof(Capture_t));
    if (!cap) {
        log_syserr("new capture error.");
        return NULL;
    }
    memset(cap, 0, sizeof(Capture_t));
    cap->raw_fd = socket_packet_new();
    cap->tp = thread_pool_new(10);
    thread_pool_set_auto_extern(cap->tp);
    thread_pool_set_limit_number(cap->tp, 2000);
    return cap;
}

int capture_run(Capture_t* cap)
{
    int len;
    pthread_t th_id;
    while (1) {
        Packet_t* packet = packet_new();
        len = recvfrom(cap->raw_fd, packet->buf, sizeof(packet->buf), 0, NULL, NULL);
        if (len < 0) {
            log_syserr("capture recvfrom error.");
            return -1;
        }
        packet->len = len;
        thread_pool_add_task(cap->tp, Analyze, (void*)packet);
        // log_trace("running number:%d", thread_pool_get_running_number(cap->tp));
        // log_trace("task number:%d", thread_pool_get_task_number(cap->tp));
    }
    return 1;
}

void capture_free(Capture_t* cap)
{
    close(cap->raw_fd);
    thread_pool_destroy(cap->tp);
    free(cap);
}